import torch
import torch.nn as nn
import json


class ResidualBlock(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(ResidualBlock, self).__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_dim, input_dim)
        self.bn2 = nn.BatchNorm1d(input_dim)

    def forward(self, x):
        identity = x
        out = self.linear1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.linear2(out)
        out = self.bn2(out)
        out += identity  # 残差连接
        out = self.relu(out)
        return out


class DynamicModelBuilder:
    @staticmethod
    def build_layer(layer_config, input_dim):
        """根据配置动态构建网络层"""
        layer_type = layer_config["type"]
        params = layer_config.get("params", {})
        activation = layer_config.get("activation", None)

        # 创建层
        if layer_type == "Linear":
            params["in_features"] = input_dim
            layer = nn.Linear(**params)
            output_dim = params["out_features"]
        elif layer_type == "Dropout":
            layer = nn.Dropout(**params)
            output_dim = input_dim
        elif layer_type == "BatchNorm1d":
            params["num_features"] = input_dim
            layer = nn.BatchNorm1d(**params)
            output_dim = input_dim
        elif layer_type == "ResidualBlock":
            # 残差块需要输入维度和隐藏层维度
            params["input_dim"] = input_dim
            layer = ResidualBlock(**params)
            output_dim = input_dim  # 残差块不改变维度
        else:
            raise ValueError(f"不支持的层类型: {layer_type}")

        # 添加激活函数
        activation_layer = None
        if activation:
            if activation == "ReLU":
                activation_layer = nn.ReLU()
            elif activation == "LeakyReLU":
                activation_layer = nn.LeakyReLU(0.01)
            elif activation == "Sigmoid":
                activation_layer = nn.Sigmoid()
            elif activation == "Tanh":
                activation_layer = nn.Tanh()
            else:
                raise ValueError(f"不支持的激活函数: {activation}")

        return layer, activation_layer, output_dim


class PricePredictor(nn.Module):
    def __init__(self, config_path, num_numeric_features):
        """
        :param config_path: 配置文件路径
        :param num_numeric_features: 数值型特征数量
        """
        super(PricePredictor, self).__init__()

        # 加载模型配置
        with open(config_path, "r") as f:
            config = json.load(f)["model"]

        input_dim = num_numeric_features

        # 动态构建网络层
        self.layers = nn.ModuleList()
        current_dim = input_dim

        for layer_config in config["layers"]:
            # 构建层
            layer, activation, output_dim = DynamicModelBuilder.build_layer(
                layer_config,
                current_dim
            )

            self.layers.append(layer)
            current_dim = output_dim

            # 添加激活函数
            if activation:
                self.layers.append(activation)

        # 打印模型结构
        print("构建模型结构:")
        for i, layer in enumerate(self.layers):
            print(f"  [{i}] {layer.__class__.__name__}")

    def forward(self, x_numeric, x_film_id=None):  # x_film_id 忽略但保留接口
        x = x_numeric  # 只使用数值特征
        for layer in self.layers:
            x = layer(x)
        return x
